from torchvision.models import resnet18
import torch
from nni.algorithms.compression.pytorch.pruning import L1FilterPruner
from nni.algorithms.compression.pytorch.pruning.apply_compression import apply_compression_results
from nni.compression.pytorch.speedup import ModelSpeedup

device = "cpu" if torch.cuda.is_available() else "gpu"

model = resnet18(True).eval().to(device)
dummy_input = torch.randn((4,3,224,224)).to(device)

# pruner = L1FilterPruner(model,config_list=[{'op_types':['Conv2d'], 'sparsity':0.5}])
# pruner.compress()
# pruner.export_model("./model", "./mask")
# pruner._unwrap_model()


apply_compression_results(model, "./mask", device)
use_mask_out = model(dummy_input)

ms = ModelSpeedup(model, dummy_input, './mask')
ms.speedup_model()
use_speed_out = model(dummy_input)

if torch.allclose(use_mask_out, use_speed_out, atol=1e-7):
    print("same")
else:
    print("different")